import sqlite3
import json
from pyecharts import WordCloud
from pyecharts import Sankey

if __name__ == '__main__':

    f_data = open('data.txt','w')
    f_links = open('links.txt','w')
    conn = sqlite3.connect('test.db')
    cursor = conn.cursor()

    # cursor.execute('create table bankuai (stock_id varchar(20), stock_name varchar(20),yyb_id varchar(20),yyb_name varchar(32),bk_id varchar(20),bk_name varchar(20))')

    # 取营业部id
    cursor.execute('select bk_name as source,yyb_name as target,count(bk_id) as value from bankuai   group by bk_id')

    values = cursor.fetchall()

    print(values)

    name = []
    re = []
    for aa in values:
        print(aa[0])
        tmp = {}
        if(aa[2]>1):
            tmp['source'] = aa[1]
            tmp['target'] = aa[0]
            tmp['value'] = aa[2]
            name.append(aa[0])
            name.append(aa[1])
            re.append(tmp)
        # print(aa)

    # print(json.dumps(re))
    print(re)

    name = list(set(name))

    names = []
    for nm in name:
        nt = {}
        nt['name'] = nm
        names.append(nt)

    # print(names)



    conn.close()

    f_data.write(json.dumps(re))
    f_links.write(json.dumps(names))

    sankey = Sankey("桑基图示例", width=1200, height=1200)
    sankey.add("sankey", names, re, line_opacity=0.2,
               line_curve=0.5, line_color='source',
               is_label_show=True, label_pos='right')
    sankey.render()
